//-*-C++-*-
/***************************************************************************
 *
 *   Copyright (C) 2025 by Will Gauvin
 *   Licensed under the Academic Free License version 2.1
 *
 ***************************************************************************/

#ifndef __dsp_RescaleMeanStdCalculator_h
#define __dsp_RescaleMeanStdCalculator_h

#include "dsp/Rescale.h"

namespace dsp
{
  /**
   * @brief a class used to calculate the mean and standard deviation of timeseries data
   */
  class RescaleMeanStdCalculator : public Rescale::ScaleOffsetCalculator {

  public:
    //! Default constructor
    RescaleMeanStdCalculator() = default;

    //! Default destructor
    virtual ~RescaleMeanStdCalculator() = default;

    /**
     * @brief initialise the calculator to allow for allocation of arrays and/or scratch space
     *
     * @param input the input timeseries of the Rescale operation, needed to get correct nchan, npol and ndim
     * @param ndat the number of samples to be used to calculate the scales and offset for
     * @param output_time_total flag that instructs the Calculator to compute an F-scrunch time series of rescaled data
     */
    void init(const dsp::TimeSeries* input, uint64_t nsample, bool output_time_total) override;

    /**
     * @brief used to sample the current input timeseries to allow calculation of scales and offsets
     *
     * This method stores accumulated values of the sample values and sample values squared across all dimensions
     * for each channel and polarisation.
     *
     * @param input the input timeseries to sample the data for, it may be either FPT or TFP ordered data.
     * @param start_dat the start time sample, may not be 0
     * @param end_dat the end time sample, may not be input->get_ndat() but may be less than that.
     * @param output_time_total an indicator of whether to accumulate the time sample values.
     * @returns the ending time sample (should be end_dat)
     */
    uint64_t sample_data(const dsp::TimeSeries* input, uint64_t start_dat, uint64_t end_dat, bool output_time_total) override;

    /**
     * @brief computes the scale and offset for each channel and polarisation of the sample data.
     *
     * @param nsample the total number of samples that Rescale has used to sample data.
     */
    void compute(uint64_t nsample) override;

    /**
     * @brief resets the sampled data
     */
    void reset_sample_data() override;

    /**
     * @brief get a pointer to the mean values for the channels of given polarisations
     */
    const double* get_mean(unsigned ipol) const override;

    /**
     * @brief get a pointer to the variance values for the channels of given polarisations
     */
    const double* get_variance(unsigned ipol) const override;

  private:
    //! sum of time samples, ordered by [pol][chan]
    std::vector< std::vector<double> > freq_total;

    //! sum of the square of the time samples, ordered by [pol][chan]
    std::vector< std::vector<double> > freq_totalsq;
  };

} // namespace dsp

#endif // __dsp_RescaleMeanStdCalculator_h
